package com.aikg.kgrag_java.repository;

import com.aikg.kgrag_java.pojo.entity.neo4j.Node;
import com.aikg.kgrag_java.pojo.po.LastOrNextNodeSearchPO;
import org.springframework.data.neo4j.repository.Neo4jRepository;
import org.springframework.data.neo4j.repository.query.Query;
import org.springframework.data.repository.query.Param;
import org.springframework.stereotype.Repository;
import org.springframework.transaction.annotation.Transactional;

import java.util.List;


@Repository
@Transactional
public interface Neo4jRetrievalRepository extends Neo4jRepository<Node, Long> {

    /**
     * 是否存在某一索引
     *
     * @param indexName: 索引的名字
     * @param nodeLabel: 索引对应的节点label名称
     * @param embedding: 索引对应的节点属性名称
     * @return 是否存在该索引
     */
    @Query("""
            SHOW INDEXES YIELD name, type, entityType, labelsOrTypes, 
            properties, options WHERE name=$indexName and labelsOrTypes[0] = $nodeLabel and properties[0] = $embeddingNodeProperty 
            RETURN count(name)>0;
           """)
    boolean hasIndex(@Param("indexName") String indexName,
                     @Param("nodeLabel") String nodeLabel,
                     @Param("embeddingNodeProperty") String embedding);

    /**
     * 创建向量索引
     *
     * @param indexName: 索引名称
     * @param nodeLabel: 索引对应的节点label名称
     * @param embedding: 索引对应的节点属性名称
     */
    @Query("""
            CALL db.index.vector.createNodeIndex(
            $indexName,
            $nodeLabel,
            $embeddingNodeProperty,
            toInteger(1024),
            'COSINE' )
            """)
    void createVectorIndex(@Param("indexName") String indexName,
                           @Param("nodeLabel") String nodeLabel,
                           @Param("embeddingNodeProperty") String embedding);

    /**
     * 相似度检索
     * @param indexName: index的名称
     * @param topK: 前 k个最相似的节点
     * @param embedding: 检索的向量
     * @param scoreThreshold: 检索的相似度大于该阈值的
     * @return 重要节点的id数组
     */
    @Query("""
            CALL db.index.vector.queryNodes($indexName, $topK, $embedding) 
            YIELD node, score 
            WHERE score >= $scoreThreshold
            return node
            """)
    List<Node> similaritySearch(@Param("indexName") String indexName,
                                   @Param("topK") int topK,
                                   @Param("embedding") List<Double> embedding,
                                   @Param("scoreThreshold") double scoreThreshold);


    /**
     * 检索重要节点的下一跳
     * @param id: 重要节点的id
     * @return 下一跳节点
     */
    @Query("""
            match (n)-[r]->(m)
            where id(n) = $id
            return r.relationName as relationName, m as targetNode 
            """)
    List<LastOrNextNodeSearchPO> nextNodeSearch(
            @Param("id") Long id
    );

    /**
     * 检索重要节点的上一跳
     * @param id: 重要节点的id
     * @return 上一跳节点
     */
    @Query("""
            match (n)-[r]->(m)
            where id(m) = $id and type(r) <> 'IMPLEMENT' 
            return r.relationName as relationName, n as headNode 
            """)
    List<LastOrNextNodeSearchPO> lastNodeSearch(
            @Param("id") Long id
    );
}
